#include "embeddingmanager.h"

EmbeddingManager::EmbeddingManager(QObject *parent) : QObject(parent)
{
    m_session = nullptr;
    m_vectorDB = nullptr;
}

bool EmbeddingManager::initialize()
{
//    qDebug() << "initialize";
    m_session = text_embedding_create_session();
    int initResult = text_embedding_init_session(m_session);
    if (initResult != 0) {
        return false;
    }
    char *model_info=nullptr;
    text_embedding_get_model_info(m_session,&model_info);
    qDebug()<<"model_info:"<<model_info;
//    EmbeddingResult *result = nullptr;
//    bool sucess = text_embedding(m_session, "我爱学习", &result);
//    qDebug()<<"sucess:"<<sucess;
//    qDebug()<<"result:"<<result;

    m_vectorDB = new VectorDatabase(this);
    if (!m_vectorDB->loadFromFile()) {
        return false;
    }
    return true;
}

void EmbeddingManager::processDocuments(const QStringList &filePaths,int maxChunkSize, int minChunkSize)
{
    for (const QString &filePath : filePaths) {

        qDebug() << "Processing file:" << filePath;
        QFile file(filePath);
        if (!file.exists()) {
            qDebug() << "File does not exist:" << filePath;
            continue;
        }
        QStringList chunks = readTextFile(filePath,maxChunkSize,minChunkSize);

        for (const QString &chunk : chunks) {

            if (m_vectorDB->documentExists(filePath, chunk)) {
                qDebug() << "already exist";
                continue;
            }
//            qDebug() << "chunk:" << chunk;
            QVector<float> embedding = getEmbeddingSync(chunk);
//            qDebug() << "embedding vector size:" << embedding.size();
            if (!embedding.isEmpty()) {
                DocumentChunk docChunk(chunk, filePath, embedding);
                m_vectorDB->insertDocument(docChunk);
                qDebug() << "Processed chunk from" << filePath;
            }
        }
    }
}

QStringList EmbeddingManager::queryRAG(const QString &question, int topK)
{
    QVector<float> questionEmbedding = getEmbeddingSync(question);
    if (questionEmbedding.isEmpty()) {
        return QStringList();
    }

    QVector<DocumentChunk> similarChunks = m_vectorDB->searchSimilar(questionEmbedding, topK);

    QStringList contexts;
    for (const DocumentChunk &chunk : similarChunks) {
        contexts.append(chunk.text);
    }

    return contexts;
}

QStringList EmbeddingManager::readTextFile(const QString &filePath,int maxChunkSize, int minChunkSize)
{
    QFile file(filePath);
    if (!file.open(QIODevice::ReadOnly | QIODevice::Text)) {
        qDebug() << "Failed to open file:" << filePath;
        return QStringList();
    }

    QTextStream in(&file);
    in.setCodec("UTF-8");
    QString content = in.readAll();
    file.close();
//    qDebug()<<content;
    return splitTextIntoChunks(content,maxChunkSize,minChunkSize);
}

QStringList EmbeddingManager::splitTextIntoChunks(const QString &text, int maxChunkSize, int minChunkSize)
{
    QStringList chunks;
    if (text.isEmpty()) return chunks;

    // 按照段落分割
    QStringList paragraphs;
    QString textCopy = text;
    QTextStream stream(&textCopy);
    QString currentPara;

    while (!stream.atEnd()) {
        QString line = stream.readLine().trimmed();
        if (line.isEmpty()) {
            if (!currentPara.isEmpty()) {
                paragraphs.append(currentPara);
                currentPara.clear();
            }
        } else {
            if (!currentPara.isEmpty()) {
                currentPara += "\n";
            }
            currentPara += line;
        }
    }

    if (!currentPara.isEmpty()) {
        paragraphs.append(currentPara);
    }

    for (const QString &para : paragraphs) {
        if (para.length() <= maxChunkSize) {
            chunks.append(para);
            continue;
        }

        // 过长段落的分割
        QStringList sentences;
        int lastPos = 0;
        for (int i = 0; i < para.length(); ++i) {
            QChar c = para[i];
            if (c == '。' || c == '！' || c == '？' || c == '\n') {
                sentences.append(para.mid(lastPos, i - lastPos + 1).trimmed());
                lastPos = i + 1;
            }
        }

        if (lastPos < para.length()) {
            sentences.append(para.mid(lastPos).trimmed());
        }

        // 构建最终块
        QString currentChunk;
        for (const QString &sentence : sentences) {
            if (currentChunk.isEmpty()) {
                currentChunk = sentence;
            }
            else if (currentChunk.length() + sentence.length() + 1 <= maxChunkSize) {
                currentChunk += "\n" + sentence;
            }
            else {
                chunks.append(currentChunk);
                currentChunk = sentence;
            }
        }

        if (!currentChunk.isEmpty()) {
            chunks.append(currentChunk);
        }
    }

    // 过滤过小的不是句子的块
    QStringList finalChunks;
    for (const QString &chunk : chunks) {
        if (chunk.length() >= minChunkSize || isCompleteSentence(chunk)) {
            finalChunks.append(chunk);
        }
    }

    return finalChunks;
}
// 辅助函数：判断是否是完整句子
bool EmbeddingManager::isCompleteSentence(const QString &text)
{
    qDebug() << "判断是否是完整句子:"<<text;
    // 简单判断：以中文句号、问号、感叹号结尾
    static QList<QChar> sentenceEnders = {'。', '！', '？', ';'};
    return !text.isEmpty() && sentenceEnders.contains(text.back());
}

QVector<float> EmbeddingManager::getEmbeddingSync(const QString &text)
{
//    qDebug() << "getEmbeddingSync";
    EmbeddingResult *result = nullptr;

    QByteArray textData = text.toUtf8();
    const char* textPtr = textData.constData();

    // 向量化
    bool success = text_embedding(m_session, textPtr, &result);
    if(!success){
        qDebug() << "Failed in text_embedding";
        return QVector<float>();
    }

    // 提取结果
    QVector<float> embeddingVector;
    if(result) {
        int len = embedding_result_get_vector_length(result);
        float* vectorData = embedding_result_get_vector_data(result);

        if(len > 0 && vectorData) {
            embeddingVector.resize(len);
            std::copy(vectorData, vectorData + len, embeddingVector.begin());
        }

        embedding_result_destroy(&result);
    }

    qDebug() << "Embedding vector size:" << embeddingVector.size();
    return embeddingVector;
}

void EmbeddingManager::resetVectorDatabase()
{
    QMutexLocker locker(&m_mutex);
    if (m_vectorDB) {
        bool success = m_vectorDB->clearDatabase();
        if (success) {
            qDebug() << "Vector database reset successfully";
        } else {
            qDebug() << "Failed to reset vector database";
        }
    }
}

